/*
 *  hooklib.cpp - part of HookLib32
 *  Copyright (C) 2012  Dustin Escoffery
 *
 *  HookLib32 is free software: you can redistribute it and/or
 *  modify it under the terms of the Lesser GNU General Public License
 *  as published by the Free Software Foundation, version 3 of the
 *  license.
 *
 *  HookLib32 is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  Lesser GNU General Public License for more details.
 *
 *  You should have received a copy of the Lesser GNU General Public
 *  License along with HookLib32.  If not, see:
 *      http://www.gnu.org/licenses/lgpl-3.0.txt
 *
 *  CONTRIBUTORS:
 *      Dustin Escoffery
 */

#include "hooklib.h"

// Gets the import descriptor for hModule
PIMAGE_IMPORT_DESCRIPTOR GetImportDescriptor(HMODULE hModule)
{
   PCHAR pImageBase =
      (PCHAR)hModule;
   PIMAGE_DOS_HEADER pDOS =
      PIMAGE_DOS_HEADER(pImageBase);
   PIMAGE_NT_HEADERS pNT =
      (PIMAGE_NT_HEADERS)(pImageBase + pDOS->e_lfanew);
   IMAGE_DATA_DIRECTORY imgDir =
      pNT->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
   DWORD dwImportDirectory =
      imgDir.VirtualAddress;
   return (PIMAGE_IMPORT_DESCRIPTOR)(pImageBase + dwImportDirectory);
}

// Hooks an import of hModule
PROC WriteImportAddress(HMODULE hModule, PCSTR pszModuleName,
   LPCSTR lpProcName, PROC pfnHooked)
{
   // Get the description of the import section
   PIMAGE_IMPORT_DESCRIPTOR pImportDesc = GetImportDescriptor(hModule);

   // Find target DLL
   while (pImportDesc->Name != 0)
   {
      PSTR pszModName = (PSTR)((PBYTE)hModule + pImportDesc->Name);
      if (!stricmp(pszModName, pszModuleName))
      {
         // Get import address table
         PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)
            ((PBYTE)hModule + pImportDesc->FirstThunk);
         // Get import name table
         PIMAGE_THUNK_DATA pOrig = (PIMAGE_THUNK_DATA)
            ((PBYTE)hModule + pImportDesc->OriginalFirstThunk);

         // Find target import
         while (pThunk->u1.Function != 0)
         {
            bool isOrd = pOrig->u1.Ordinal & IMAGE_ORDINAL_FLAG;
            PIMAGE_IMPORT_BY_NAME pImportName = (PIMAGE_IMPORT_BY_NAME)
               ((PBYTE)hModule + pOrig->u1.AddressOfData);

            if ((isOrd && LOWORD(lpProcName) == LOWORD(pOrig->u1.Ordinal)) ||
                (!isOrd && !strcmp(lpProcName, (LPCSTR)pImportName->Name)))
            {
               PROC* ppfn = (PROC*)&pThunk->u1.Function;
               PROC pfnOriginal = *ppfn;

               // Set writable permission
               MEMORY_BASIC_INFORMATION mbi;
               VirtualQuery(ppfn, &mbi,
                  sizeof(MEMORY_BASIC_INFORMATION));
               VirtualProtect(mbi.BaseAddress, mbi.RegionSize,
                  PAGE_READWRITE, &mbi.Protect);

               // Hook the function
               *ppfn = *pfnHooked;

               // Restore permissions
               DWORD dwOldProtect;
               VirtualProtect(mbi.BaseAddress, mbi.RegionSize,
                  mbi.Protect, &dwOldProtect);

               return pfnOriginal;
            }
            pThunk++;
            pOrig++;
         }
      }
      pImportDesc++;
   }
   return NULL;
}
